use std::{path::PathBuf, process::Stdio, time::Duration};

use anyhow::Context;
use clap::Parser as _;
use tokio::process::Command;

mod manifest;
mod proc;

#[derive(Debug, clap::Parser)]
#[command(version, about, long_about = None)]
struct Args {
    #[arg(long, help = "Path to cuttlefish directory")]
    cf_path: PathBuf,
    #[arg(long, help = "Path to MITM server directory")]
    mitm_path: PathBuf,
    #[arg(long, help = "Path to DroidBot directory")]
    droidbot_path: PathBuf,
    #[arg(long, help = "Path to JADX directory")]
    jadx_path: PathBuf,
    #[arg(long, help = "Index of instance (cvd-N) to attach to")]
    instance_num: u8,
    #[arg(short, long, help = "Path to input APK file")]
    input: PathBuf,
    #[arg(short, long, help = "Path to output directory")]
    output: PathBuf,
    #[arg(long, default_value_t = false, help = "Skip powerwash")]
    no_powerwash: bool,
    #[arg(long, default_value_t = false, help = "Disable DroidBot")]
    manual: bool,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    env_logger::Builder::new()
        .filter_level(log::LevelFilter::Info)
        .parse_default_env()
        .init();

    let args: Args = Args::parse();
    let cf_home = args.cf_path;
    let instance_num = args.instance_num;
    let adb_port = 6519 + instance_num as u16; // instance_num starts from 1, adb port starts from 6520
    let tcp_mitm_port = 16519 + instance_num as u16;
    let mtap_interface = format!("cvd-mtap-{:02}", instance_num);

    let adb_serial = format!("0.0.0.0:{}", adb_port);
    let adb_path = cf_home.join("bin/adb");

    let apk_path = args.input;
    if !apk_path.exists() {
        anyhow::bail!("Input file {:?} does not exist", apk_path);
    }

    tokio::fs::remove_dir_all(&args.output).await.ok();
    tokio::fs::create_dir_all(&args.output).await?;

    log::info!("Parsing APK...");

    proc::run_command(
        Command::new(args.jadx_path.join("bin/jadx"))
            .arg("--no-src")
            .arg("--output-dir-res")
            .arg(args.output.join("jadx-res"))
            .arg(&apk_path),
        Duration::from_secs(60),
    )
    .await
    .context("Running JADX")?;

    let manifest_path = args.output.join("jadx-res/AndroidManifest.xml");
    if !manifest_path.exists() {
        anyhow::bail!("Failed to decode AndroidManifest.xml");
    }
    log::info!("Successfully parsed APK, got AndroidManifest.xml");

    let package_name = match manifest::parse_manifest(&manifest_path) {
        Ok(pi) => pi.package_name,
        Err(e) => {
            log::error!("Failed to parse manifest: {:?}", e);
            None
        }
    };

    let pcap_path = args.output.join("network.pcap");
    let mitm_dump_path = args.output.join("tcp.mitm");
    let keylog_path = args.output.join("keylog.txt");
    let mitm_json_path = args.output.join("mitm.json");

    if !args.no_powerwash {
        log::info!(
            "Starting to powerwash cvd-{}, ADB serial is {}",
            args.instance_num,
            adb_serial
        );

        proc::run_command(
            Command::new(cf_home.join("bin/powerwash_cvd"))
                .current_dir(&cf_home)
                .env("HOME", &cf_home)
                .arg(format!("--instance_num={}", instance_num)),
            Duration::from_secs(60),
        )
        .await
        .context("Running powerwash_cvd")?;
    }

    log::info!("Waiting for ADB {} to become available", adb_serial);

    proc::run_command(
        Command::new(&adb_path)
            .arg("-s")
            .arg(&adb_serial)
            .arg("wait-for-device"),
        Duration::from_secs(60),
    )
    .await
    .context("Running adb wait-for-device")?;

    log::info!("Device successfully restarted");

    let mut mitmweb_child = Command::new(args.mitm_path.join(".venv/bin/mitmweb"))
        .env("SSLKEYLOGFILE", &keylog_path)
        .env("PYTHONUNBUFFERED", "1")
        .args([
            "--set",
            "ssl_insecure=true",
            "--mode",
            "transparent",
            "--showhost",
            "--no-web-open-browser",
        ])
        .arg("--listen-port")
        .arg(format!("{}", tcp_mitm_port))
        .arg("--set")
        .arg(format!(
            "confdir={}",
            args.mitm_path.join("config").display()
        ))
        .arg("--set")
        .arg(format!("output_json_path={}", mitm_json_path.display()))
        .arg("-s")
        .arg(args.mitm_path.join("mitm_addon.py"))
        .arg("-w")
        .arg(&mitm_dump_path)
        .kill_on_drop(true)
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()?;

    proc::start_forwarding_log(&mut mitmweb_child, "mitmweb");

    // sudo setcap cap_net_raw,cap_net_admin=eip /usr/bin/tcpdump
    let mut tcpdump_child = Command::new("tcpdump")
        .arg("-i")
        .arg(&mtap_interface)
        .arg("-s")
        .arg("65535")
        .arg("-w")
        .arg(&pcap_path)
        .kill_on_drop(true)
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()?;

    proc::start_forwarding_log(&mut tcpdump_child, "tcpdump");

    let ctrl_c = tokio::signal::ctrl_c();

    if args.manual {
        let package_name = if let Some(pn) = package_name {
            pn
        } else {
            anyhow::bail!("Could not launch in manual mode without package name");
        };

        log::info!("Installing APK...");

        proc::run_command(
            Command::new(&adb_path)
                .arg("-s")
                .arg(&adb_serial)
                .arg("install")
                .arg("-r")
                .arg("-g")
                .arg(&apk_path),
            Duration::from_secs(120),
        )
        .await
        .context("Running adb install")?;

        log::info!("Launching app...");

        proc::run_command(
            Command::new(&adb_path)
                .arg("-s")
                .arg(&adb_serial)
                .arg("shell")
                .arg("monkey")
                .arg("-p")
                .arg(&package_name)
                .arg("1"),
            Duration::from_secs(30),
        )
        .await
        .context("Running adb shell monkey")?;

        log::info!("App launched, feel free to explore");

        tokio::select! {
            _ = mitmweb_child.wait() => {
                log::info!("mitmweb has stopped");
            }
            _ = tcpdump_child.wait() => {
                log::info!("tcpdump has stopped");
            }
            _ = ctrl_c => {
                log::info!("Got ctrl-c, stopping");
            }
        }

        log::info!("Uninstalling app...");

        let del_res = proc::run_command(
            Command::new(&adb_path)
                .arg("-s")
                .arg(&adb_serial)
                .args(["shell", "pm", "uninstall"])
                .arg(&package_name),
            Duration::from_secs(120),
        )
        .await;
        if let Err(e) = del_res {
            log::warn!("Failed to uninstall app: {:?}", e);
        }
    } else {
        let mut droidbot_child = Command::new(args.droidbot_path.join("venv/bin/python"))
            .env("PATH", cf_home.join("bin")) // Use ADB in Cuttlefish
            .env("PYTHONUNBUFFERED", "1")
            .arg(args.droidbot_path.join("start.py"))
            .arg("-a")
            .arg(apk_path)
            .arg("-o")
            .arg(args.output.join("droidbot"))
            .arg("-d")
            .arg(&adb_serial)
            .arg("-keep_env")
            .arg("-grant_perm")
            .kill_on_drop(true)
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;

        proc::start_forwarding_log(&mut droidbot_child, "droidbot");

        tokio::select! {
            _ = mitmweb_child.wait() => {
                log::info!("mitmweb has stopped");
            }
            _ = droidbot_child.wait() => {
                log::info!("DroidBot has stopped");
            }
            _ = tcpdump_child.wait() => {
                log::info!("tcpdump has stopped");
            }
            _ = ctrl_c => {
                log::info!("Got ctrl-c, stopping");
            }
        }

        log::info!("Stopping DroidBot");
        proc::sigint_and_wait(&mut droidbot_child).await;
    }

    log::info!("Stopping mitmweb");
    proc::sigint_and_wait(&mut mitmweb_child).await;
    log::info!("Stopping tcpdump");
    proc::sigint_and_wait(&mut tcpdump_child).await;

    Ok(())
}
